package com.yfs.cloud.filter;

import io.netty.buffer.ByteBufAllocator;

import org.apache.commons.lang3.StringEscapeUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.server.ServerWebExchange;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.validation.constraints.NotEmpty;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Optional;

/**
* @Company: 影法師工作室
* @Description: XSS过滤器 
* @author 影法師  
* @date 2020年8月28日 下午1:54:00 
* @version V1.0
 */
//@Component // 好烦, 这个 xss 过滤搞不好
@ConfigurationProperties("config.xss")
@SuppressWarnings("all")
public class XssFilter implements GlobalFilter, Ordered {

	private Logger logger = LoggerFactory.getLogger(this.getClass());

	private List<XssWhiteUrl> whiteUrls;

	@Override
	public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
		ServerHttpRequest request = exchange.getRequest();
		URI uri = request.getURI();
		String method = request.getMethodValue();
		if (this.white(uri.getPath(), method)) {
			return chain.filter(exchange);
		}
		if ((HttpMethod.POST.name().equals(method) || HttpMethod.PUT.name().equals(method))) {
			return DataBufferUtils.join(request.getBody()).flatMap(d -> Mono.just(Optional.of(d)))
					.defaultIfEmpty(Optional.empty()).flatMap(optional -> {
						// 取出body中的参数
						String bodyString = "";
						if (optional.isPresent()) {
							byte[] oldBytes = new byte[optional.get().readableByteCount()];
							optional.get().read(oldBytes);
							bodyString = new String(oldBytes, StandardCharsets.UTF_8);
						}
						HttpHeaders httpHeaders = request.getHeaders();
						// 执行XSS清理
						logger.info("{} - [{}] XSS处理前参数:{}", method, uri.getPath(), bodyString);
						try {
							bodyString = StringEscapeUtils.escapeJson(bodyString).replace("\\r\\n", "").replace("\\",
									"");
						} catch (Exception e) {
							System.out.println(e);
							logger.error("json 解析异常  : {}", bodyString);
							throw new RuntimeException();
						}
						logger.info("{} - [{}] XSS处理后参数：{}", method, uri.getPath(), bodyString);
						ServerHttpRequest newRequest = request.mutate().uri(uri).build();

						// 重新构造body
						byte[] newBytes = bodyString.getBytes(StandardCharsets.UTF_8);
						DataBuffer bodyDataBuffer = toDataBuffer(newBytes);
						Flux<DataBuffer> bodyFlux = Flux.just(bodyDataBuffer);

						// 重新构造header
						HttpHeaders headers = new HttpHeaders();
						headers.putAll(httpHeaders);
						// 由于修改了传递参数，需要重新设置CONTENT_LENGTH，长度是字节长度，不是字符串长度
						int length = newBytes.length;
						headers.remove(HttpHeaders.CONTENT_LENGTH);
						headers.setContentLength(length);
						headers.set(HttpHeaders.CONTENT_TYPE, "application/json;charset=utf8");
						// 重写ServerHttpRequestDecorator，修改了body和header，重写getBody和getHeaders方法
						newRequest = new ServerHttpRequestDecorator(newRequest) {
							@Override
							public Flux<DataBuffer> getBody() {
								return bodyFlux;
							}

							@Override
							public HttpHeaders getHeaders() {
								return headers;
							}
						};

						return chain.filter(exchange.mutate().request(newRequest).build());
					});
		} else {
			return chain.filter(exchange);
		}
	}

	/**
	 * 是否是白名单
	 *
	 * @param url    路由
	 * @param method 请求方式
	 * @return true/false
	 */
	private boolean white(String url, String method) {
		return whiteUrls != null && whiteUrls.contains(XssWhiteUrl.builder().url(url).method(method).build());
	}

	/**
	 * 字节数组转DataBuffer
	 *
	 * @param bytes 字节数组
	 * @return DataBuffer
	 */
	private DataBuffer toDataBuffer(byte[] bytes) {
		NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);
		DataBuffer buffer = nettyDataBufferFactory.allocateBuffer(bytes.length);
		buffer.write(bytes);
		return buffer;
	}

	/**定义的大一点, 作为最后一道屏障*/
	public static final int ORDER = 10;

	@Override
	public int getOrder() {
		return ORDER;
	}

	@Validated
	private static class XssWhiteUrl {

		@NotEmpty
		private String url;

		@NotEmpty
		private String method;

		public static XssWhiteUrl builder() {
			return new XssWhiteUrl();
		}

		public XssWhiteUrl url(String url) {
			this.url = url;
			return this;
		}

		public XssWhiteUrl method(String method) {
			this.method = method;
			return this;
		}

		public XssWhiteUrl build() {
			return this;
		}
	}

	public static void main(String[] args) {
		String str = "abcd\\r\\ndeflsg\\r\\n \"id\": 2020,\\r\\n    \"username\": \"yfscfs\"\\r\\n ";
		System.out.println(str.replace("\\r\\n", ""));
	}
}
